# gpt 添加相关注释与代码优化
from cmath import atan
from typing import Tuple
import cv2
import numpy as np
from matplotlib import pyplot as plt

class RoadDetector:
    """道路与楼梯检测器
    功能：
    - 基于颜色阈值的道路中线检测
    - 楼梯位置检测
    - 环形道路/分叉路处理
    - 实时跟踪数据显示
    
    坐标系说明：
    (0,0)—————→ x轴
    |
    |
    ↓ 
    y轴
    """

    # 滑动平均参数（用于数据平滑）
    __avgK = 0       # 中线斜率滑动平均值
    __avgB = 0       # 中线截距滑动平均值
    __pushedCount = 0# 参与平均的帧数计数器
    __roundIn = False # 环形道路标志位
    __splitTick = 0   # 分叉路持续时间计数
    __maskPos = None  # ROI区域坐标 (x,y)
    __useRight = False# 分叉时使用右侧路径标志

    # 图像处理相关
    img_proced = None # 预处理后的图像
    img_mask = None   # 最终的二值掩模
    __stairMode = False # 楼梯检测模式标志
    __stairDetect = False # 楼梯检测激活标志

    def __init__(self, roadThreshold, stairThreshold, backThreshold, maskKSize):
        """初始化检测器
        Args:
            roadThreshold: 道路颜色阈值列表 [([H_min,S_min,V_min], [H_max,S_max,V_max])]
            stairThreshold: 楼梯颜色阈值 (lower, upper)
            backThreshold: 背景排除阈值 (lower, upper)
            maskKSize: 形态学操作核尺寸（奇数）
        """
        self.__roadThres = roadThreshold    # 道路颜色阈值组
        self.__stairThres = stairThreshold  # 楼梯颜色阈值
        self.__maskKSize = maskKSize        # 形态学核大小
        self.__backThres = backThreshold    # 背景过滤阈值

    def setMaskSize(self, maskSize):
        """动态设置形态学核尺寸（用于适应不同分辨率）"""
        self.__maskKSize = maskSize

    def __detectRoad(self, mask, img, holeClip):
        """道路中线检测核心算法
        Args:
            mask: 二值化道路掩模
            img: 原始图像（用于可视化调试）
            holeClip: 是否启用空洞处理
        Returns:
            tuple: (直线斜率k, 截距b) 或 None
        
        算法流程：
        1. 纵向扫描每一行，检测有效像素区域
        2. 处理道路中间的空洞（如斑马线造成的断裂）
        3. 最小二乘法拟合最佳直线
        """
        midY = []  # 存储每行中点的x坐标
        midX = []  # 对应的y坐标
        holeEnd = False  # 空洞处理完成标志

        # 纵向扫描处理（从下往上）
        for y in range(mask.shape[0]):
            # 提取当前行有效像素列坐标
            valid_pixels = np.where(mask[y, :] == 255)[0]
            
            if valid_pixels.size > 80:  # 有效像素数阈值
                # 空洞检测与处理
                if holeClip and not holeEnd:
                    right = valid_pixels.max()
                    left = valid_pixels.min()
                    # 检测潜在空洞（实际宽度 >> 有效像素数）
                    if (right - left) - valid_pixels.size > 5:
                        # 定位空洞边界
                        black_zones = np.where(mask[y, :] == 0)[0]
                        in_black = black_zones[black_zones >= left]
                        
                        if in_black.size >= 80:  # 确认显著空洞
                            # 根据标志选择左右路径
                            if self.__useRight:
                                split_pos = black_zones[in_black.argmax()]
                                valid_right = valid_pixels[valid_pixels > split_pos]
                                if valid_right.size > 60:
                                    valid_pixels = valid_right
                            else:
                                split_pos = black_zones[in_black.argmin()]
                                valid_left = valid_pixels[valid_pixels < split_pos]
                                if valid_left.size > 60:
                                    valid_pixels = valid_left
                        else:
                            holeEnd = True  # 结束空洞处理
                
                # 记录当前行中点
                midX.append(y)
                midY.append(int(valid_pixels.mean()))
                
                # 可视化调试标记
                if len(midX) >= 2 and img is not None:
                    color = (0,255,0) if self.__stairMode else (255,0,0)
                    cv2.line(img, (midY[-1], midX[-1]), (midY[-2], midX[-1]), color, 1)

        # 最小二乘法拟合直线 y = kx + b
        if len(midY) < 20:  # 有效数据不足
            return None
        
        # 构建矩阵方程
        y = np.array(midY)
        x = np.array(midX)
        A = np.asmatrix([[len(x), x.sum()], 
                   [x.sum(), (x**2).sum()]])
        b_vec = np.array([y.sum(), (x*y).sum()])
        
        try:
            # 求解线性方程组
            det = np.linalg.det(A)
            if abs(det) < 1e-4:  # 矩阵奇异情况
                return None
            solution = np.linalg.solve(A, b_vec)
            b, k = solution[0], solution[1]
            
            # 坐标系转换（ROI到原始图像）
            b += self.__maskPos[0] - self.__maskPos[1] * k
            return (k, b)
        except np.linalg.LinAlgError:
            return None

    def __imagePreprocess(self, img, maxSel):
        """图像预处理流水线
        Args:
            img: 输入图像（HSV格式）
            maxSel: 是否选择最大连通区域
        Returns:
            处理后的二值化掩模
        
        处理步骤：
        1. 颜色阈值处理
        2. 背景排除
        3. 形态学开闭运算
        4. 最大连通区域选择
        """
        # 颜色阈值处理
        mask = np.zeros_like(img[:, :, 0])
        if self.__stairMode:
            lower,upper = self.stairThres
            mask = cv2.inRange(img, np.array(lower),np.array(upper))
        else:
            for threshold in self.__roadThres:
                lower,upper = threshold
                temp_mask = cv2.inRange(img,np.array(lower),np.array(upper))
                mask = cv2.bitwise_or(mask, temp_mask)
        
        # 背景排除处理
        lower,upper = self.__backThres
        neg_mask = cv2.inRange(img,np.array(lower),np.array(upper))
        mask = cv2.bitwise_xor(mask, neg_mask)
        
        # 形态学处理（去噪+填充）
        kernel = np.ones((self.__maskKSize, self.__maskKSize), np.uint8)
        big_k = np.ones((2*self.__maskKSize, 2*self.__maskKSize), np.uint8) # 添加by lbx

        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)  # 去除小噪ss
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) # 填充小孔洞
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, big_k)  # 添加by lbx
        
        # 最大连通区域选择
        if maxSel:
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
            if contours:
                max_contour = max(contours, key=cv2.contourArea)
                cont_mask = np.zeros_like(mask)
                cv2.drawContours(cont_mask, [max_contour], -1, 255, cv2.FILLED)
                
                # 计算有效区域边界
                x, y, w, h = cv2.boundingRect(cont_mask)
                if h / img.shape[0] > 0.9:  # 有效区域高度阈值
                    self.__maskPos = (x, y)
                    return cont_mask[y:y+h, x:x+w]
        return mask

    def __detectStair(self, mask, image):
        """楼梯位置检测
        通过计算最大连通域的质心确定楼梯位置
        """
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        if not contours:
            return (0, 0)
        
        # 寻找最大轮廓并计算质心
        max_contour = max(contours, key=cv2.contourArea)
        M = cv2.moments(max_contour)
        if M["m00"] == 0:
            return (0, 0)
        
        cx = int(M["m10"] / M["m00"])
        cy = int(M["m01"] / M["m00"])
        cv2.circle(image, (cx, cy), 5, (0,255,0), 3)
        return (cx, cy)

    def __judgeTop(self, image):
        """判断是否到达楼梯顶部
        通过分析纵向密度变化检测平台区域
        """
        # 轮廓筛选与密度分析
        contours, _ = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        density = np.zeros(image.shape[0])
        
        for y in range(image.shape[0]):
            density[y] = np.sum(image[y, :]) // 255
        
        # 使用滑动平均检测密度突变
        kernel = np.ones(5)/5
        smooth_density = np.convolve(density, kernel, mode='valid')
        derivatives = np.diff(smooth_density)
        
        # 寻找显著上升沿
        rise_edges = np.where(derivatives > 5)[0]
        return rise_edges.size == 0  # 无突变表示到达平台

    def update(self, image):
        """主更新函数（逐帧调用）
        执行模式判断与数据更新
        """
        # 图像预处理
        self.img_proced = self.__imagePreprocess(image, not self.__stairMode)
        self.img_mask = self.img_proced.copy()
        # 模式分支处理
        if not self.__stairMode:
            # 道路中线检测
            result = self.__detectRoad(self.img_proced, image, not self.__stairMode)
            if result is None: 
                self.__pushedCount = 0
                print("failure in decect road.")
                return
            k, b = result
            # print(f'k:{k},b{b}')  # 右拐k = -1.35 b = 360
            # 滑动平均滤波
            # print(f'original avgK {}')
            self.__avgK = (self.__avgK * self.__pushedCount + k) / (self.__pushedCount + 1)
            # print(f'self._avgK: {self.__avgK}')
            # print(f'self._avgb: {self.__avgB}')
            self.__avgB = (self.__avgB * self.__pushedCount + b) / (self.__pushedCount + 1)
        else:
            # 楼梯质心检测
            x, y = self.__detectStair(self.img_proced, image)
            self.__avgK = (self.__avgK * self.__pushedCount + y) / (self.__pushedCount + 1)
            self.__avgB = (self.__avgB * self.__pushedCount + x) / (self.__pushedCount + 1)
        
        self.__pushedCount += 1

        # 楼梯顶部检测逻辑
        if self.__stairMode and self.__stairDetect:
            if self.__judgeTop(self.img_proced):
                self.__stairTopCount = min(self.__stairTopCount + 1, 50)
                if self.__stairTopCount >= 50:
                    self.__stairMode = False  # 退出楼梯模式
            else:
                self.__stairTopCount = max(self.__stairTopCount - 1, 0)

    def showTracking(self, image):
        """可视化跟踪结果"""
        if not self.__stairMode:
            # 绘制道路中线
            pt1 = (int(self.__avgB), 0)
            pt2 = (int(pt1[0] + image.shape[0] * self.__avgK), image.shape[0])
            cv2.line(image, pt1, pt2, (255, 0, 0), 2)
            # 标记ROI区域
            cv2.rectangle(image, self.__maskPos, 
                         (self.__maskPos[0]+self.img_proced.shape[1], 
                          self.__maskPos[1]+self.img_proced.shape[0]), 
                         (0,255,0), 1)
        else:
            # 标记楼梯位置
            cv2.circle(image, (int(self.__avgB), int(self.__avgK)), 10, (0,255,0), 3)

    def getTrackingData(self):
        """获取跟踪数据
        Returns:
            道路模式: (k, b, 起点, 终点)
            楼梯模式: (x, y)
        """
        if self.__pushedCount == 0 or self.img_proced is None:
            return None
        
        self.__pushedCount = 0  # 重置计数器
        if self.__stairMode:
            return (self.__avgB, self.__avgK)
        else:
            pt1 = (int(self.__avgB), 0)
            pt2 = (int(pt1[0] + self.img_proced.shape[0] * self.__avgK), 
                  self.img_proced.shape[0])
            # print(f'pt1:{pt1} pt2:{pt2}')
            # print(f'self__avgK:{self.__avgK}')
            return (self.__avgK, self.__avgB, pt1, pt2) # Pt1 在上,Pt2在下

    def setStairMode(self, enable):
        """切换楼梯检测模式
        Args:
            enable: True进入楼梯模式，False返回道路模式
        """
        if self.__stairMode != enable:
            self.__pushedCount = 0  # 模式切换时重置历史数据
        self.__stairMode = enable
